package com.study.one.pulsar;

import org.apache.pulsar.client.api.HashingScheme;
import org.apache.pulsar.client.api.Message;
import org.apache.pulsar.client.api.TopicMetadata;
import org.apache.pulsar.client.impl.RoundRobinPartitionMessageRouterImpl;
import org.apache.pulsar.shade.com.google.common.annotations.VisibleForTesting;

import java.time.Clock;
import java.util.concurrent.ThreadLocalRandom;


public class PulsarRouter extends RoundRobinPartitionMessageRouterImpl {

    private static final PulsarRouter INSTANCE = new PulsarRouter();

    public PulsarRouter() {
        this(Math.abs(ThreadLocalRandom.current().nextInt()), Clock.systemUTC());
    }

    @VisibleForTesting
    public PulsarRouter(int startPtnIdx, Clock clock) {
        super(
                HashingScheme.Murmur3_32Hash,
                startPtnIdx,
                true,
                1,
                clock);
    }

    public static PulsarRouter of() {
        return INSTANCE;
    }

    @Override
    public int choosePartition(Message msg, TopicMetadata metadata) {
        if (msg.hasKey() || msg.getSequenceId() < 0) {
            return super.choosePartition(msg, metadata);
        }
        return (int) (msg.getSequenceId() % metadata.numPartitions());
    }

}
